from option import Option
import torch
import random
import numpy as np

from dataset.dataloader import ImageDataLoader
seed = 7
torch.manual_seed(seed) # cpu
torch.cuda.manual_seed(seed) # gpu
np.random.seed(seed) # numpy
random.seed(seed) # random and transform
from trainer import *
from model.Ulike import total_model
from checkpoint import CheckPoint

os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DIVICES'] = "" # gpu序号
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True # cudnn


def main(opt=None):
    start_time = time.time()

    opt = opt or Option()

    checkpoint = CheckPoint(opt=opt)


    # create dataloader
    data_loader = ImageDataLoader(opt.train_image_path,opt.test_image_path,opt.image_path,opt.train_mask_path,opt.test_mask_path,
                                  opt.train_batch_size,opt.test_batch_size,opt.n_threads)
    train_loader, test_loader = data_loader.get_dataloader()
    print("==>Finish loading data\n")

    model = total_model()

    model_dict = torch.load(opt.check_point).state_dict()
    model.load_state_dict(model_dict)
    print(model)

    Ulike_model = utils.dataparallel(model=model, ngpus=opt.n_gpus, gpu0=opt.gpu0)
    print("==>Finish loading model\n")
    start_epoch = 0

    # create trainer
    trainer = Trainer(model=model, opt=opt)

    for epoch in range(start_epoch, opt.nepoch):
        print("==>Epoch ", epoch)
        train_loss = trainer.train(train_loader=train_loader,epoch=epoch)
        test_loss = trainer.evaluate(test_loader=test_loader,epoch=epoch)

        checkpoint.save_path(epoch=epoch, model=trainer.model, adam_opts = trainer.adam_optimizer)

    end_time = time.time()
    time_interval = end_time - start_time
    print("==>Time is: %f\n"
          % (time_interval))
    log_str = "Time: %f\n" % (time_interval)
    utils.writelog(log_str, opt.log_file)
if __name__ == '__main__':
    main()